import torch




a_edge, a_value = (torch.LongTensor(
[
        [0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7, 8, 8, 8, 9, 9],
        [5, 7, 8, 5, 8, 6, 7, 9, 6, 8, 6, 9, 0, 1, 2, 3, 4, 0, 2, 0, 1, 3, 2, 4]]
),torch.FloatTensor(([0.1000, 0.1000, 0.3000, 0.1000, 0.3000, 0.1000, 0.1000, 0.3000, 0.1000,
        0.3000, 0.1000, 0.3000, 0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.2000,
        0.2000, 0.4000, 0.4000, 0.4000, 0.4000, 0.4000])))

b_edge, b_value = (torch.LongTensor(
[[0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7, 8, 8, 8, 9, 9],
        [5, 7, 8, 5, 8, 6, 7, 9, 6, 8, 6, 9, 0, 1, 2, 3, 4, 0, 2, 0, 1, 3, 2, 4]]
),torch.FloatTensor(([0.2000, 0.2000, 0.3000, 0.2000, 0.3000, 0.2000, 0.2000, 0.3000, 0.2000,
        0.3000, 0.2000, 0.3000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000,
        0.1000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000])))
num_nodes =  10

# 定义了全局邻接矩阵的稀疏形式。
mat_a = torch.sparse_coo_tensor(a_edge, a_value, (num_nodes, num_nodes)).to(a_edge.device)
mat_b = torch.sparse_coo_tensor(b_edge, b_value, (num_nodes, num_nodes)).to(a_edge.device)
print(mat_a)
print(mat_b)
print(mat_a.shape)
print(mat_b.shape)
print("####################")
print(torch.sparse.mm(mat_a, mat_b))
# 非零元素的连续存储形式，提高运算效率
mat = torch.sparse.mm(mat_a, mat_b).coalesce()
print(mat)
edges, values = mat.indices(), mat.values()
print("a_edge",a_edge)
print("b_edge",b_edge)
print("edges",edges)
